from flask import Blueprint, render_template, request, jsonify
from app.models.neural_network import FaultDiagnosisModel
from app.utils.data_processor import DataProcessor
import os
from app.config import Config
import json

# 创建蓝图并定义路由为'/training'
training_bp = Blueprint('training', __name__, url_prefix='/training')
model = FaultDiagnosisModel()
data_processor = DataProcessor()


@training_bp.route('/', methods=['GET'])
def index():
    """训练页面"""
    # 获取可用的数据集
    datasets = []
    if os.path.exists(Config.DATA_FOLDER):
        for file in os.listdir(Config.DATA_FOLDER):
            if file.endswith(('.csv', '.xlsx', '.xls')):
                datasets.append(file)

    return render_template('training.html', datasets=datasets)


@training_bp.route('/train', methods=['POST'])
def train():
    """开始训练"""
    try:
        data = request.get_json()
        dataset = data.get('dataset')
        target_column = data.get('target_column')
        epochs = int(data.get('epochs', 10))

        if not dataset or not target_column:
            return jsonify({
                'status': 'error',
                'message': '缺少必要参数'
            }), 400

        # 加载数据
        file_path = os.path.join(Config.DATA_FOLDER, dataset)
        df = data_processor.load_data(file_path)

        # 预处理数据
        train_loader, test_loader = data_processor.preprocess_data(df, target_column)

        # 训练模型
        history = model.train_model(train_loader, epochs=epochs)

        # 评估模型
        accuracy, _, _ = model.evaluate(test_loader)

        # 保存模型
        model.save_model()

        return jsonify({
            'status': 'success',
            'message': '训练完成',
            'accuracy': accuracy,
            'history': history
        })

    except Exception as e:
        return jsonify({
            'status': 'error',
            'message': str(e)
        }), 400


@training_bp.route('/upload', methods=['POST'])
def upload_data():
    """上传数据集"""
    if 'file' not in request.files:
        return jsonify({
            'status': 'error',
            'message': '没有文件被上传'
        }), 400

    file = request.files['file']
    if file.filename == '':
        return jsonify({
            'status': 'error',
            'message': '没有选择文件'
        }), 400

    if file and '.' in file.filename and \
       file.filename.rsplit('.', 1)[1].lower() in Config.ALLOWED_EXTENSIONS:
        filename = file.filename
        file.save(os.path.join(Config.DATA_FOLDER, filename))
        return jsonify({
            'status': 'success',
            'message': '文件上传成功',
            'filename': filename
        })
    else:
        return jsonify({
            'status': 'error',
            'message': '不支持的文件类型'
        }), 400
